import React, { useEffect, useMemo, useState } from 'react';

interface TreeProps {
  id: string;
  children?: string;
}

interface ManualControlProps {
  onAfterDataLoad?: () => void;
  treeProps: TreeProps;
}

const useManualTreeExpandControl = (props: ManualControlProps) => {
  const idField = props?.treeProps?.id || 'id';
  const childrenField = props?.treeProps.children || 'children';

  /** 当前是否已经全部展开, 用来控制全部折叠/展开 */
  const [isExpandAll, setIsExpandAll] = useState(false);
  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
  const [autoExpandParent, setAutoExpandParent] = useState(true);

  /** 需要暂存一下 treeData */
  const [treeData, setTreeData] = useState<any[]>([]);
  // useMemo 计算一下 listData
  const listData = useMemo(() => {
    const _dataList: any[] = [];
    // 将 treeData 中的数据放到 listData 里面, 如果 有children 则 children 也放到 listData 里面
    const loop = (data: any[]) => {
      data.forEach((item) => {
        _dataList.push({ ...item });
        if (item[childrenField]) {
          loop(item[childrenField]);
        }
      });
    };
    loop(treeData);

    return _dataList;
  }, [treeData]);

  useEffect(() => {
    props?.onAfterDataLoad?.();
  }, [listData]);

  /** 暂存一下 treeData 数据用于后面的计算 */
  const cacheTreeData = (data: any[]) => {
    setTreeData(data);
  };

  /** 用于接管树的展开事件 */
  const onExpand = (newExpandedKeys: React.Key[]) => {
    setExpandedKeys(newExpandedKeys);
    setAutoExpandParent(false);
  };

  /** 用于接管表格的展开事件 */
  const onTableExpand = (expanded: boolean, record: any) => {
    if (expanded) {
      // 展开, 需要往expandedKeys中添加
      onExpand([...expandedKeys, record[idField]]);
    } else {
      // 收起, 需要从expandedKeys中删除
      onExpand(expandedKeys.filter((key) => key !== record[idField]));
    }
  };

  const onExpandAutoParent = (newExpandedKeys: React.Key[]) => {
    setExpandedKeys(newExpandedKeys);
    setAutoExpandParent(true);
  };

  /** 展开全部 */
  const expandAll = () => {
    const keys = listData.map((item) => item[idField]);
    setExpandedKeys(keys);
    setIsExpandAll(true);
    setAutoExpandParent(true);
  };

  /** 折叠全部 */
  const collapseAll = () => {
    setExpandedKeys([]);
    setAutoExpandParent(false);
    setIsExpandAll(false);
  };

  /** 切换展开/折叠全部 */
  const toggleExpandAll = () => {
    if (isExpandAll) {
      collapseAll();
    } else {
      expandAll();
    }
  };

  return {
    treeData,
    listData,
    isExpandAll,
    expandedKeys,
    autoExpandParent,
    cacheTreeData,
    onExpand,
    onExpandAutoParent,
    expandAll,
    collapseAll,
    toggleExpandAll,
    onTableExpand,
  };
};

export default useManualTreeExpandControl;
